"""
This is the main file of the KlingAI package.
Doc: https://docs.qingque.cn/d/home/eZQClW07IFEuX1csc-VejdY2M?identityId=1oEG9JKKMFv#section=h.9mwdken9otn8
"""

import httpx
import jwt
import time

from common import setup_logger, TASK_STATUS
from config import PropertiesReader


class KlingAI:
    def __init__(self):
        properties_reader = PropertiesReader()
        self.base_path = 'https://api.klingai.com'
        self.log = setup_logger('klingai')

        self.image_to_video_url = f'{self.base_path}/v1/videos/image2video'
        self.ak = properties_reader.get_property('ak')
        self.sk = properties_reader.get_property('sk')

    def get_query_image_to_video_task_status_url(self, id):
        return f'{self.base_path}/v1/videos/image2video/{id}'

    def encode_jwt_token(self):
        headers = {
            "alg": "HS256",
            "typ": "JWT"
        }
        payload = {
            "iss": self.ak,
            "exp": int(time.time()) + 1800,  # 有效时间，此处示例代表当前时间+1800s(30min)
            "nbf": int(time.time()) - 5  # 开始生效的时间，此处示例代表当前时间-5秒
        }

        token = jwt.encode(payload, self.sk, headers=headers)
        return token

    def build_headers(self):
        token = self.encode_jwt_token()
        return {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {token}'
        }

    def image_to_video(self,
                       model_name: str = 'kling-v1', # kling-v1, kling-v1-5
                       mode: str = 'pro', # 枚举值：std(标准模式（标准），基础模式，性价比高)，pro(专家模式（高品质），高表现模式，生成视频质量更佳)
                       duration: int = 5, # 单位：s秒，枚举5，10,
                       image: str = 'https://h2.inkwai.com/bs2/upload-ylab-stunt/se/ai_portal_queue_mmu_image_upscale_aiweb/3214b798-e1b4-4b00-b7af-72b5b0417420_raw_image_0.jpg',
                       prompt: str = '宇航员站起身走了',
                       cfg_scale: float = 0.5,
                       negative_prompt: str = None,
                       static_mask: str = 'https://h2.inkwai.com/bs2/upload-ylab-stunt/ai_portal/1732888177/cOLNrShrSO/static_mask.png',
                       dynamic_masks: dict = [{
                           "mask": "https://h2.inkwai.com/bs2/upload-ylab-stunt/ai_portal/1732888130/WU8spl23dA/dynamic_mask_1.png",
                           "trajectories": [
                               {"x": 279, "y": 219}, {"x": 417, "y": 65}
                           ]
                       }]
                       ):
        client = httpx.Client()
        headers = self.build_headers()

        try:
            data = {
                "model_name": model_name,
                "mode": mode,
                "duration": str(duration),
                "image": image,
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "cfg_scale": cfg_scale,
                "static_mask": static_mask,
                "dynamic_masks": dynamic_masks
            }
            response = client.post(self.image_to_video_url, json=data, headers=headers)
            response.raise_for_status()
            resp = response.json()
            self.log.info(resp)
            return resp
        except httpx.HTTPStatusError as exc:
            self.log.error(f"HTTP error occurred: {exc}")
        except Exception as exc:
            self.log.error(f"An error occurred: {exc}")
        finally:
            client.close()

    def query_task_status(self, id):
        try:
            client = httpx.Client()
            headers = self.build_headers()
            url = self.get_query_image_to_video_task_status_url(id)
            response = client.get(url, headers=headers)
            response.raise_for_status()  # 如果响应状态码不是200，将抛出异常

            status_response = response.json()
            self.log.info(status_response)

            videos_array = []
            if status_response['message'] == 'SUCCEED':
                self.log.info('SUCCEED')
                data_node = status_response['data']
                task_id = data_node['task_id']
                task_status = data_node['task_status']
                task_result_node = data_node['task_result']

                if task_status == 'processing':
                    self.log.info(f'task in processing, task_id: {task_id}')
                    return TASK_STATUS.PROCESSING, []
                videos_array = task_result_node['videos']
                self.log.info(f'task_id:    {task_id}')
                self.log.info(f'task_status:    {task_status}')
            return TASK_STATUS.SUCCESS, videos_array
        except httpx.HTTPStatusError as exc:
            self.log.error(f"HTTP error occurred: {exc}")
        except Exception as exc:
            self.log.error(f"An error occurred: {exc}")
        finally:
            client.close()
